#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MCP连接管理器

解决MCP服务连接稳定性问题，提供自动重连和健康检查功能。

作者: MCP决策优化团队
版本: 1.0.0
日期: 2024
"""

import os
import sys
import time
import json
import subprocess
import threading
from pathlib import Path
from typing import Optional, Dict, Any

class MCPConnectionManager:
    """MCP连接管理器"""
    
    def __init__(self, config_path: str = "mcp-config.json"):
        self.config_path = Path(config_path)
        self.project_root = Path(__file__).parent
        self.process: Optional[subprocess.Popen] = None
        self.is_running = False
        self.health_check_interval = 5  # 秒
        self.max_restart_attempts = 3
        self.restart_count = 0
        
        # 加载配置
        self.config = self._load_config()
        
    def _load_config(self) -> Dict[str, Any]:
        """加载MCP配置"""
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"[警告] 加载配置失败: {e}")
            return self._get_default_config()
    
    def _get_default_config(self) -> Dict[str, Any]:
        """获取默认配置"""
        return {
            "mcpServers": {
                "decision-optimization": {
                    "command": "python",
                    "args": ["src/main.py"],
                    "cwd": str(self.project_root),
                    "env": {
                        "PYTHONPATH": "src"
                    },
                    "timeout": 30,
                    "restart": True
                }
            }
        }
    
    def start_service(self) -> bool:
        """启动MCP服务"""
        if self.is_running:
            print("[INFO] MCP服务已在运行中")
            return True
            
        try:
            server_config = self.config["mcpServers"]["decision-optimization"]
            
            # 设置环境变量
            env = os.environ.copy()
            if "env" in server_config:
                env.update(server_config["env"])
            
            # 启动进程
            cmd = [server_config["command"]] + server_config["args"]
            cwd = server_config.get("cwd", str(self.project_root))
            
            print(f"[启动] MCP服务: {' '.join(cmd)}")
            print(f"[目录] 工作目录: {cwd}")
            
            self.process = subprocess.Popen(
                cmd,
                cwd=cwd,
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1,
                universal_newlines=True
            )
            
            # 等待服务启动
            time.sleep(2)
            
            if self.process.poll() is None:
                self.is_running = True
                self.restart_count = 0
                print(f"[成功] MCP服务启动成功 (PID: {self.process.pid})")
                
                # 启动健康检查
                self._start_health_check()
                return True
            else:
                stdout, stderr = self.process.communicate()
                print(f"[失败] MCP服务启动失败")
                print(f"stdout: {stdout}")
                print(f"stderr: {stderr}")
                return False
                
        except Exception as e:
            print(f"[错误] 启动MCP服务时发生错误: {e}")
            return False
    
    def stop_service(self):
        """停止MCP服务"""
        if self.process and self.is_running:
            try:
                self.process.terminate()
                self.process.wait(timeout=5)
                print("[成功] MCP服务已停止")
            except subprocess.TimeoutExpired:
                self.process.kill()
                print("[强制] 强制终止MCP服务")
            except Exception as e:
                print(f"[警告] 停止服务时发生错误: {e}")
            finally:
                self.is_running = False
                self.process = None
    
    def restart_service(self) -> bool:
        """重启MCP服务"""
        if self.restart_count >= self.max_restart_attempts:
            print(f"[失败] 已达到最大重启次数 ({self.max_restart_attempts})，停止重启")
            return False
            
        print(f"[重启] 重启MCP服务 (第 {self.restart_count + 1} 次)")
        self.stop_service()
        time.sleep(1)
        
        self.restart_count += 1
        return self.start_service()
    
    def check_health(self) -> bool:
        """检查服务健康状态"""
        if not self.process:
            return False
            
        # 检查进程是否还在运行
        if self.process.poll() is not None:
            print("[警告] MCP服务进程已退出")
            return False
            
        return True
    
    def _start_health_check(self):
        """启动健康检查线程"""
        def health_check_loop():
            while self.is_running:
                time.sleep(self.health_check_interval)
                
                if not self.check_health():
                    print("[检测] 检测到服务异常，尝试重启...")
                    if not self.restart_service():
                        print("[失败] 服务重启失败，停止监控")
                        break
        
        health_thread = threading.Thread(target=health_check_loop, daemon=True)
        health_thread.start()
    
    def get_status(self) -> Dict[str, Any]:
        """获取服务状态"""
        return {
            "is_running": self.is_running,
            "process_id": self.process.pid if self.process else None,
            "restart_count": self.restart_count,
            "max_restart_attempts": self.max_restart_attempts,
            "health_check_interval": self.health_check_interval
        }

def main():
    """主函数"""
    # 设置控制台编码为UTF-8
    import sys
    if sys.platform == "win32":
        import os
        os.system('chcp 65001 > nul')
    
    print("MCP连接管理器")
    print("=" * 50)
    
    manager = MCPConnectionManager()
    
    try:
        # 启动服务
        if manager.start_service():
            print("\n[提示] 服务已启动，按 Ctrl+C 停止服务")
            
            # 保持运行
            while True:
                time.sleep(1)
                
    except KeyboardInterrupt:
        print("\n[停止] 收到停止信号")
    finally:
        manager.stop_service()
        print("[退出] 程序退出")

if __name__ == "__main__":
    main()